# -*- coding: utf-8 -*-
"""
SHAP分析脚本

使用SHAP（SHapley Additive exPlanations）解释模型预测结果，帮助理解不同特征对预测的影响。
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import pickle
import json
import shap
from model import load_model
from sklearn.preprocessing import OneHotEncoder

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 创建SHAP分析结果目录
if not os.path.exists('shap_results'):
    os.makedirs('shap_results')


class TorchModelWrapper:
    """
    PyTorch模型包装器，用于SHAP分析
    
    参数:
        model (nn.Module): PyTorch模型
    """
    def __init__(self, model):
        self.model = model
        self.model.eval()
    
    def __call__(self, X):
        """
        调用模型进行预测
        
        参数:
            X (np.ndarray): 输入特征
            
        返回:
            np.ndarray: 预测结果
        """
        # 转换为PyTorch张量
        X_tensor = torch.FloatTensor(X)
        
        # 使用模型预测
        with torch.no_grad():
            predictions = self.model(X_tensor).numpy()
        
        return predictions


def load_data_and_model():
    """
    加载数据和模型
    
    返回:
        tuple: (X_train, y_train, X_test, y_test, model, feature_names, original_feature_names, encoders)
    """
    print("加载数据和模型...")
    
    # 从.npy文件加载处理后的数据
    X_train = np.load('output/X_train.npy')
    X_test = np.load('output/X_test.npy')
    y_train = np.load('output/y_train.npy')
    y_test = np.load('output/y_test.npy')
    
    # 加载特征名称
    with open('output/feature_names.pkl', 'rb') as f:
        feature_names = pickle.load(f)
    
    # 设置原始特征名称
    original_feature_names = ['高分子基体（缩写）', '无机添加剂（缩写）', '其余材料', '加工方法']
    
    # 加载编码器
    with open('output/encoders.pkl', 'rb') as f:
        encoders = pickle.load(f)
    
    # 加载最佳模型信息
    with open('output/best_model_info.json', 'r') as f:
        best_model_info = json.load(f)
    
    best_model_name = best_model_info['name']
    print(f"最佳模型: {best_model_name}")
    
    # 加载模型配置
    model_configs = {
        'mlp_model': {
            'type': 'mlp',
            'params': {
                'hidden_dims': [64, 32],
                'dropout_rate': 0.2,
                'use_batch_norm': True,
                'activation': 'relu'
            }
        },
        'residual_model': {
            'type': 'residual',
            'params': {
                'hidden_dims': [64, 128, 64],
                'num_res_blocks': 2,
                'dropout_rate': 0.2,
                'use_batch_norm': True
            }
        },
        'dual_branch_model': {
            'type': 'dual_branch',
            'params': {
                'hidden_dims': [64, 128],
                'branch_dims': [64, 32],
                'dropout_rate': 0.2,
                'use_batch_norm': True,
                'num_heads': 4
            }
        },
        'ensemble_model': {
            'type': 'ensemble',
            'params': {
                'num_models': 3,
                'model_type': 'dual_branch',
                'hidden_dims': [64, 128],
                'dropout_rate': 0.2,
                'use_batch_norm': True
            }
        }
    }
    
    # 获取模型配置
    model_config = model_configs[best_model_name]
    
    # 加载模型
    input_dim = len(feature_names)
    model = load_model(
        model_config['type'],
        input_dim,
        f"models/{best_model_name}.pth",
        **model_config['params']
    )
    
    # 设置为评估模式
    model.eval()
    
    return X_train, y_train, X_test, y_test, model, feature_names, original_feature_names, encoders


def create_feature_names_mapping(feature_names, original_feature_names, encoders):
    """
    创建特征名称映射
    
    参数:
        feature_names (list): 特征名称列表
        original_feature_names (list): 原始特征名称列表
        encoders (dict): 特征编码器字典
        
    返回:
        list: 映射后的特征名称列表
    """
    # 创建映射后的特征名称列表
    mapped_feature_names = []
    
    # 处理每个原始特征
    feature_idx = 0
    for col in original_feature_names:
        # 检查是否有编码器
        if col in encoders:
            # 获取类别名称
            categories = encoders[col].categories_[0]
            
            # 为每个类别创建特征名称
            for category in categories:
                mapped_feature_names.append(f"{col}_{category}")
            
            feature_idx += len(categories)
        else:
            # 对于连续特征，直接使用原始名称
            mapped_feature_names.append(col)
            feature_idx += 1
    
    return mapped_feature_names


def run_shap_analysis(model, X_train, X_test, feature_names):
    """
    运行SHAP分析
    
    参数:
        model (nn.Module): 训练好的神经网络模型
        X_train (np.ndarray): 训练集特征
        X_test (np.ndarray): 测试集特征
        feature_names (list): 特征名称列表
        
    返回:
        tuple: (explainer, shap_values_test)
    """
    print("运行SHAP分析...")
    
    # 创建模型包装器
    model_wrapper = TorchModelWrapper(model)
    
    # 创建SHAP解释器
    # 使用KernelExplainer，适用于任何模型
    # 使用训练集的子集作为背景数据
    n_background = min(100, X_train.shape[0])  # 最多使用100个样本作为背景
    background_data = X_train[:n_background]
    
    explainer = shap.KernelExplainer(model_wrapper, background_data)
    
    # 计算测试集的SHAP值
    # 使用测试集的子集以提高计算效率
    n_samples = min(50, X_test.shape[0])  # 最多使用50个样本
    X_test_sample = X_test[:n_samples]
    
    # 计算SHAP值（对于多输出模型，返回一个列表，每个元素对应一个输出）
    shap_values_test = explainer.shap_values(X_test_sample)
    
    return explainer, shap_values_test


def visualize_shap_values(explainer, shap_values_test, X_test, feature_names):
    """
    可视化SHAP值
    
    参数:
        explainer (shap.Explainer): SHAP解释器
        shap_values_test (list): 测试集的SHAP值
        X_test (np.ndarray): 测试集特征
        feature_names (list): 特征名称列表
    """
    print("可视化SHAP值...")
    
    # 使用测试集的子集
    n_samples = min(50, X_test.shape[0])
    X_test_sample = X_test[:n_samples]
    
    # 对于多输出模型，分别可视化每个输出
    output_names = ['太阳反射率', '红外发射率']
    
    for i, output_name in enumerate(output_names):
        print(f"可视化 {output_name} 的SHAP值...")
        
        # 摘要图：显示每个特征的平均绝对SHAP值
        plt.figure(figsize=(12, 8))
        shap.summary_plot(
            shap_values_test[i],
            X_test_sample,
            feature_names=feature_names,
            show=False
        )
        plt.title(f"{output_name} - 特征重要性摘要")
        plt.tight_layout()
        plt.savefig(f"shap_results/{output_name}_summary_plot.png", dpi=300, bbox_inches='tight')
        plt.close()
        
        # 条形图：显示每个特征的平均绝对SHAP值
        plt.figure(figsize=(12, 10))
        shap.summary_plot(
            shap_values_test[i],
            X_test_sample,
            feature_names=feature_names,
            plot_type="bar",
            show=False
        )
        plt.title(f"{output_name} - 特征重要性条形图")
        plt.tight_layout()
        plt.savefig(f"shap_results/{output_name}_bar_plot.png", dpi=300, bbox_inches='tight')
        plt.close()
        
        # 依赖图：显示特征值与SHAP值的关系
        # 选择前5个最重要的特征
        mean_abs_shap = np.abs(shap_values_test[i]).mean(axis=0)
        top_indices = np.argsort(mean_abs_shap)[-5:]
        
        for idx in top_indices:
            feature_name = feature_names[idx]
            plt.figure(figsize=(10, 6))
            shap.dependence_plot(
                idx,
                shap_values_test[i],
                X_test_sample,
                feature_names=feature_names,
                show=False
            )
            plt.title(f"{output_name} - {feature_name} 依赖图")
            plt.tight_layout()
            plt.savefig(f"shap_results/{output_name}_{feature_name}_dependence_plot.png", dpi=300, bbox_inches='tight')
            plt.close()
        
        # 力图：显示单个预测的解释
        # 选择前3个样本进行可视化
        for j in range(min(3, X_test_sample.shape[0])):
            plt.figure(figsize=(12, 6))
            shap.force_plot(
                explainer.expected_value[i],
                shap_values_test[i][j],
                X_test_sample[j],
                feature_names=feature_names,
                matplotlib=True,
                show=False
            )
            plt.title(f"{output_name} - 样本 {j+1} 的力图")
            plt.tight_layout()
            plt.savefig(f"shap_results/{output_name}_sample_{j+1}_force_plot.png", dpi=300, bbox_inches='tight')
            plt.close()


def analyze_feature_importance(shap_values_test, feature_names):
    """
    分析特征重要性
    
    参数:
        shap_values_test (list): 测试集的SHAP值
        feature_names (list): 特征名称列表
        
    返回:
        dict: 特征重要性分析结果
    """
    print("分析特征重要性...")
    
    # 对于多输出模型，分别分析每个输出
    output_names = ['太阳反射率', '红外发射率']
    
    # 创建结果字典
    results = {}
    
    for i, output_name in enumerate(output_names):
        # 计算每个特征的平均绝对SHAP值
        mean_abs_shap = np.abs(shap_values_test[i]).mean(axis=0)
        
        # 创建特征重要性DataFrame
        importance_df = pd.DataFrame({
            '特征': feature_names,
            '重要性': mean_abs_shap
        })
        
        # 按重要性排序
        importance_df = importance_df.sort_values('重要性', ascending=False)
        
        # 保存到CSV
        importance_df.to_csv(f"shap_results/{output_name}_feature_importance.csv", index=False, encoding='utf-8-sig')
        
        # 添加到结果字典
        results[output_name] = importance_df.to_dict('records')
        
        # 打印前10个最重要的特征
        print(f"\n{output_name} - 前10个最重要的特征:")
        for _, row in importance_df.head(10).iterrows():
            print(f"{row['特征']}: {row['重要性']:.4f}")
    
    return results


def analyze_feature_interactions(explainer, X_test, feature_names):
    """
    分析特征交互
    
    参数:
        explainer (shap.Explainer): SHAP解释器
        X_test (np.ndarray): 测试集特征
        feature_names (list): 特征名称列表
    """
    print("分析特征交互...")
    
    # 使用测试集的子集
    n_samples = min(50, X_test.shape[0])
    X_test_sample = X_test[:n_samples]
    
    # 对于多输出模型，分别分析每个输出
    output_names = ['太阳反射率', '红外发射率']
    
    try:
        for i, output_name in enumerate(output_names):
            print(f"分析 {output_name} 的特征交互...")
            
            # 计算SHAP交互值
            # 注意：这可能非常耗时，特别是对于大型模型和数据集
            shap_interaction_values = explainer.shap_interaction_values(X_test_sample)
            
            # 可视化交互矩阵
            plt.figure(figsize=(12, 10))
            
            # 计算平均绝对交互值
            mean_abs_interaction = np.abs(shap_interaction_values[i]).mean(axis=0)
            
            # 创建热图
            plt.imshow(mean_abs_interaction, cmap='viridis')
            plt.colorbar(label='平均绝对SHAP交互值')
            
            # 设置坐标轴标签
            plt.xticks(np.arange(len(feature_names)), feature_names, rotation=90)
            plt.yticks(np.arange(len(feature_names)), feature_names)
            
            plt.title(f"{output_name} - 特征交互矩阵")
            plt.tight_layout()
            plt.savefig(f"shap_results/{output_name}_interaction_matrix.png", dpi=300, bbox_inches='tight')
            plt.close()
            
            # 分析最强的交互
            # 获取非对角线元素
            interaction_matrix = mean_abs_interaction.copy()
            np.fill_diagonal(interaction_matrix, 0)  # 将对角线元素设为0
            
            # 找到最强的交互
            n_top = 10  # 前10个最强的交互
            flat_indices = np.argsort(interaction_matrix.flatten())[-n_top:]
            row_indices, col_indices = np.unravel_index(flat_indices, interaction_matrix.shape)
            
            # 创建交互强度DataFrame
            interactions = []
            for r, c in zip(row_indices, col_indices):
                # 确保不重复记录交互（因为矩阵是对称的）
                if r < c:
                    interactions.append({
                        '特征1': feature_names[r],
                        '特征2': feature_names[c],
                        '交互强度': interaction_matrix[r, c]
                    })
            
            # 创建DataFrame并保存
            if interactions:
                interaction_df = pd.DataFrame(interactions)
                interaction_df = interaction_df.sort_values('交互强度', ascending=False)
                interaction_df.to_csv(f"shap_results/{output_name}_top_interactions.csv", index=False, encoding='utf-8-sig')
                
                # 打印前5个最强的交互
                print(f"\n{output_name} - 前5个最强的特征交互:")
                for _, row in interaction_df.head(5).iterrows():
                    print(f"{row['特征1']} 与 {row['特征2']}: {row['交互强度']:.4f}")
    except Exception as e:
        print(f"分析特征交互时出错: {e}")
        print("跳过特征交互分析...")


def analyze_specific_samples(model, X_test, y_test, feature_names, explainer, shap_values_test):
    """
    分析特定样本
    
    参数:
        model (nn.Module): 训练好的神经网络模型
        X_test (np.ndarray): 测试集特征
        y_test (np.ndarray): 测试集标签
        feature_names (list): 特征名称列表
        explainer (shap.Explainer): SHAP解释器
        shap_values_test (list): 测试集的SHAP值
    """
    print("分析特定样本...")
    
    # 使用模型预测测试集
    model.eval()
    with torch.no_grad():
        y_pred = model(torch.FloatTensor(X_test)).numpy()
    
    # 计算每个样本的预测误差
    errors = np.abs(y_pred - y_test)
    
    # 计算每个样本的总误差（两个输出的误差之和）
    total_errors = errors.sum(axis=1)
    
    # 找到误差最大和最小的样本
    worst_idx = np.argmax(total_errors)
    best_idx = np.argmin(total_errors)
    
    # 分析这些样本
    sample_indices = [worst_idx, best_idx]
    sample_names = ['最差预测', '最佳预测']
    
    for idx, name in zip(sample_indices, sample_names):
        print(f"\n分析 {name} 样本 (索引: {idx})...")
        
        # 获取样本数据
        X_sample = X_test[idx:idx+1]
        y_true = y_test[idx]
        y_predicted = y_pred[idx]
        
        # 打印真实值和预测值
        print(f"真实值: 太阳反射率 = {y_true[0]:.4f}, 红外发射率 = {y_true[1]:.4f}")
        print(f"预测值: 太阳反射率 = {y_predicted[0]:.4f}, 红外发射率 = {y_predicted[1]:.4f}")
        print(f"误差: 太阳反射率 = {abs(y_true[0] - y_predicted[0]):.4f}, 红外发射率 = {abs(y_true[1] - y_predicted[1]):.4f}")
        
        # 为每个输出创建力图
        output_names = ['太阳反射率', '红外发射率']
        
        for i, output_name in enumerate(output_names):
            plt.figure(figsize=(12, 6))
            shap.force_plot(
                explainer.expected_value[i],
                shap_values_test[i][idx % len(shap_values_test[i])],  # 确保索引在范围内
                X_sample[0],
                feature_names=feature_names,
                matplotlib=True,
                show=False
            )
            plt.title(f"{name} - {output_name} 的力图")
            plt.tight_layout()
            plt.savefig(f"shap_results/{name}_{output_name}_force_plot.png", dpi=300, bbox_inches='tight')
            plt.close()
        
        # 创建瀑布图
        for i, output_name in enumerate(output_names):
            try:
                plt.figure(figsize=(12, 8))
                shap.plots.waterfall(
                    shap.Explanation(
                        values=shap_values_test[i][idx % len(shap_values_test[i])],
                        base_values=explainer.expected_value[i],
                        data=X_sample[0],
                        feature_names=feature_names
                    ),
                    show=False
                )
                plt.title(f"{name} - {output_name} 的瀑布图")
                plt.tight_layout()
                plt.savefig(f"shap_results/{name}_{output_name}_waterfall_plot.png", dpi=300, bbox_inches='tight')
                plt.close()
            except Exception as e:
                print(f"创建瀑布图时出错: {e}")


def generate_shap_report(feature_importance_results):
    """
    生成SHAP分析报告
    
    参数:
        feature_importance_results (dict): 特征重要性分析结果
    """
    print("生成SHAP分析报告...")
    
    # 创建报告内容
    report = []
    
    # 添加标题
    report.append("# SHAP分析报告")
    report.append("")
    report.append("## 概述")
    report.append("")
    report.append("本报告使用SHAP（SHapley Additive exPlanations）方法分析了模型的预测结果，")
    report.append("帮助理解不同特征对预测太阳反射率和红外发射率的影响。")
    report.append("")
    
    # 添加特征重要性部分
    report.append("## 特征重要性分析")
    report.append("")
    
    output_names = ['太阳反射率', '红外发射率']
    
    for output_name in output_names:
        report.append(f"### {output_name}的特征重要性")
        report.append("")
        report.append(f"![{output_name}特征重要性](../shap_results/{output_name}_bar_plot.png)")
        report.append("")
        report.append(f"前10个最重要的特征:")
        report.append("")
        report.append("| 特征 | 重要性 |")
        report.append("| --- | --- |")
        
        # 添加前10个最重要的特征
        for feature in feature_importance_results[output_name][:10]:
            report.append(f"| {feature['特征']} | {feature['重要性']:.4f} |")
        
        report.append("")
    
    # 添加特征依赖图部分
    report.append("## 特征依赖分析")
    report.append("")
    report.append("特征依赖图显示了特征值与SHAP值之间的关系，帮助理解特征如何影响预测结果。")
    report.append("")
    
    for output_name in output_names:
        report.append(f"### {output_name}的特征依赖图")
        report.append("")
        
        # 假设我们为每个输出的前5个最重要特征创建了依赖图
        top_features = [feature['特征'] for feature in feature_importance_results[output_name][:5]]
        
        for feature in top_features:
            # 替换特殊字符，以便在文件名中使用
            feature_filename = feature.replace('/', '_').replace('\\', '_').replace(' ', '_')
            report.append(f"#### {feature}")
            report.append("")
            report.append(f"![{output_name} - {feature}依赖图](../shap_results/{output_name}_{feature_filename}_dependence_plot.png)")
            report.append("")
    
    # 添加样本分析部分
    report.append("## 样本分析")
    report.append("")
    report.append("以下是对最佳预测和最差预测样本的分析，帮助理解模型在不同情况下的表现。")
    report.append("")
    
    sample_names = ['最佳预测', '最差预测']
    
    for name in sample_names:
        report.append(f"### {name}")
        report.append("")
        
        for output_name in output_names:
            report.append(f"#### {output_name}")
            report.append("")
            report.append(f"![{name} - {output_name}力图](../shap_results/{name}_{output_name}_force_plot.png)")
            report.append("")
            
            # 添加瀑布图（如果存在）
            waterfall_path = f"shap_results/{name}_{output_name}_waterfall_plot.png"
            if os.path.exists(waterfall_path):
                report.append(f"![{name} - {output_name}瀑布图](../shap_results/{name}_{output_name}_waterfall_plot.png)")
                report.append("")
    
    # 添加结论部分
    report.append("## 结论与建议")
    report.append("")
    report.append("根据SHAP分析结果，我们可以得出以下结论：")
    report.append("")
    report.append("1. 对太阳反射率影响最大的特征是...（根据实际结果填写）")
    report.append("2. 对红外发射率影响最大的特征是...（根据实际结果填写）")
    report.append("3. 特征之间的交互关系表明...（根据实际结果填写）")
    report.append("")
    report.append("基于以上分析，我们建议：")
    report.append("")
    report.append("1. 在材料设计中优先考虑...（根据实际结果填写）")
    report.append("2. 进一步研究...（根据实际结果填写）")
    report.append("3. 改进模型的...（根据实际结果填写）")
    
    # 将报告内容写入文件
    with open('shap_results/shap_analysis_report.md', 'w', encoding='utf-8') as f:
        f.write('\n'.join(report))
    
    print("SHAP分析报告已生成: shap_results/shap_analysis_report.md")


def main():
    # 加载数据和模型
    X_train, y_train, X_test, y_test, model, feature_names, original_feature_names, encoders = load_data_and_model()
    
    # 创建特征名称映射
    mapped_feature_names = create_feature_names_mapping(feature_names, original_feature_names, encoders)
    
    # 运行SHAP分析
    explainer, shap_values_test = run_shap_analysis(model, X_train, X_test, mapped_feature_names)
    
    # 可视化SHAP值
    visualize_shap_values(explainer, shap_values_test, X_test, mapped_feature_names)
    
    # 分析特征重要性
    feature_importance_results = analyze_feature_importance(shap_values_test, mapped_feature_names)
    
    # 分析特征交互
    analyze_feature_interactions(explainer, X_test, mapped_feature_names)
    
    # 分析特定样本
    analyze_specific_samples(model, X_test, y_test, mapped_feature_names, explainer, shap_values_test)
    
    # 生成SHAP分析报告
    generate_shap_report(feature_importance_results)
    
    print("\nSHAP分析完成！")
    print(f"结果已保存到 shap_results/ 目录")


if __name__ == "__main__":
    main()